##
## Copyright (C) 2009-2010 Luciano Xumerle <luciano.xumerle@gmail.com>
##
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.
##

import os, re, string, time, mmap
import sys
reload(sys)
sys.setdefaultencoding("latin1")
import bisect

try:
    import pysqlite2.dbapi2 as sqlite
except:
    import sqlite3 as sqlite


###
def getBinaryPositions( lista, rstart, rstop, fullIntersect, returnFalseStart ):
   """An iterative version of the binary search algorithm to find features range.
      Inputs are an ordered list of feature ranges and the coords to map.
      Return a list of index."""
   res=[]
   sstart=int(rstart)
   sstop=int(rstop)
   gene=[]
   if len(lista)==0:
       return gene
   start=0
   stop=len(lista)-1

   while True:
       mid=start+((stop-start)/2)
       if (stop-start)<5:
           break
       if lista[mid][1]<sstart and start<mid:
           start=mid
           continue
       if lista[mid][0]>=sstop and stop>mid:
           stop=mid+1
           continue
       break

   ## the start on the list element is less than my start
   start=stop
   while start>0 and lista[start-1][1]+1>=sstart:
       start-=1

   if returnFalseStart:
       return [start]

   while True:
       if (start > (len(lista) - 1)) or (lista[start][0]>sstop):
           break
       cur=start
       start+=1
       addFeat=False
       if fullIntersect:
           if lista[cur][0]<=sstart and sstop<=lista[cur][1]:
               addFeat=True
       else:
           if sstart>=lista[cur][0] and sstart<=lista[cur][1]:
               addFeat=True
           else:
               if sstop>=lista[cur][0] and sstop<=lista[cur][1]:
                   addFeat=True

       if addFeat:
           res.append(cur)

   res.sort()
   return res


def binaryFeatSearch( lista, rstart, rstop, fullIntersect ):
   """An iterative version of the binary search algorithm to find features range.
      Inputs are an ordered list of feature ranges and the coords to map.
      gene|start|stop|length|exonNum|."""
   res=getBinaryPositions( lista, rstart, rstop, fullIntersect, False )
   gene=[]
   for cur in res:
       if len(lista[cur])==5:
           gene.append( [lista[cur][2], lista[cur][0], lista[cur][1], lista[cur][1]-lista[cur][0], lista[cur][3], lista[cur][4]] )
       else:
           gene.append( [lista[cur][2], lista[cur][0], lista[cur][1], '=', '=', '='] )
   return gene


def createChromosomeSequence( chromosome ):
   """Converts the chromosome.fa file to a file without fasta header and spaces."""
   f = open( chromosome + '.bin', 'w' )
   for line in open( chromosome ):
       if '>' not in line:
           f.write(  line.rstrip() )
   f.close()


def getSequenceChromosome( chrom, start, length ):
   """Returns nucleotide sequence from a given chromosome."""
   st=int(start)
   ln=int(length)
   prefix=''
   while st<0:
       st+=1
       ln-=1
       prefix+='.'
   seq=''
   if st <= os.path.getsize( chrom ):
       chromFile = open(chrom, 'r+b')
       mymap = mmap.mmap(chromFile.fileno(), 0)
       mymap.seek( st )
       seq = mymap.read( ln )
       chromFile.close()
   seq=prefix + seq
   while len(seq) < int(length):
       seq+='.'
   return seq


def catSequences( fragments ):
   """Returns the sequence string.
      fragments is a list of elements [ gene, chfile, sense, start, stop]
      start positions counted from 0 and stop from 1"""
   outseq=""
   oldgene=""
   oldchrom=""
   oldsense=""
   sequence=""
   for ii in fragments:
       if oldgene != ii[0] or oldsense != ii[2] or oldchrom != ii[1]:
           if len(sequence)>0:
               if oldsense == "-":
                   sequence = reverseSeq( sequence )
               outseq += sequence
               sequence=""
           oldgene=ii[0]
           oldsense=ii[2]
           oldchrom=ii[1]
       ss = ii[3]
       ee = ii[4]
       sequence += getSequenceChromosome( ii[1], ss, ee-ss )
   if len(sequence)>0:
       if oldsense == "-":
           sequence = reverseSeq( sequence )
       outseq += sequence
   return outseq


def reverseSeq(sequence):
   """Reverse a sequence."""
   #base pairs
   RevDict={'A':'T','T':'A', 'C':'G', 'G':'C', 'N':'N'}
   reverse = ""
   #reverse the sequene
   for i in range(len(sequence)):
       reverse = RevDict[sequence[i].upper()]+reverse
   return reverse


def mergeSortedRegions( regionsDict, sortedList ):
   """"""
   dest=dict()
   # copy regionsDict to dest
   for chrom in regionsDict:
       dest[chrom]=regionsDict[chrom]

   # insert all regions
   for e in sortedList:
       ch = e[0]
       start=int(e[1])
       end=int(e[2])
       ## add chr list if missing and add value
       if ch not in dest:
           dest[ch]=[]
       dest[ch].append( [ ch, start, end ] )

   ## sort all dest lists
   for chrom in dest:
       temp=dest[chrom]
       temp.sort()
       dest[chrom]=[]
       for e in temp:
           if len(dest[chrom])==0:
               dest[chrom].append( e )
           else:
               if e[1] > dest[chrom][-1][2]+1:
                   dest[chrom].append( e )
               else:
                   if e[2] > dest[chrom][-1][2]:
                       dest[chrom][-1][2]=e[2]
   return dest


### CLASS KNOWNGENE ###
### CLASS KNOWNGENE ###
### CLASS KNOWNGENE ###
"""
The knowngene class is used to map short reads into genome features.
Each feature is defined by its range (start and stop points in the DNA sequence).
The class uses the knowngene.txt file to get each annotated feature.
Features are: exons; genes; ovrlapped gene regions; neargene range
(a user defined interval at the start and end of the gene).
"""
class knowngene:
   """The knowngene class is used to map short reads into genome features."""

   # verbose - use plog
   printLog=False

   # database name
   kgdb=''

   # db connection
   sql=''
   cur=''

   ## CACHED DATABASE ##
   splices=dict()
   cds=dict()
   utr=dict()
   exons=dict()
   genes=dict()
   neargenes=dict()
   geneLength=dict()
   geneOrient=dict()
   geneExons=dict()
   chrList=[]


   def __init__(self, kgDataBase):
       """Class Constructo.
          kgDataBase is an empty sqlite3 database
          or a previously populated (buildKGdatabase) database."""
       if kgDataBase != '':
           self.kgdb=kgDataBase
           self.sql=sqlite.connect(self.kgdb)
           self.sql.execute("PRAGMA CACHE_SIZE = 5000000")
           self.cur = self.sql.cursor()

   ### LOG METHODS ###
   ### LOG METHODS ###
   ### LOG METHODS ###
   def plog( self, msg ):
       """Prints the log message."""
       if self.printLog:
           print >>sys.stderr, time.ctime() + ' | ' + msg

   def printLogMessage(self, tf):
       """Set the class to print (tf=True) the log message or not (tf=False)"""
       self.printLog=tf

   ### CACHED QUERY ###
   ### CACHED QUERY ###
   ### CACHED QUERY ###

   def createCache(self):
       """Creates the cache needed to get fast query"""
       ## cache splices gene name
       self.cur.execute("SELECT DISTINCT id, annot FROM annotations;")
       for e in self.cur:
           if e[1][0:2]=="rs":
               continue
           if e[0] in self.splices :
               if e[1] not in self.splices[e[0]]:
                   self.splices[e[0]].append(e[1])
           else:
              self.splices[e[0]]=[]
              self.splices[e[0]].append(e[0])
              self.splices[e[0]].append(e[1])
       self.plog( "Splices genes: " + str(len(self.splices)))

       ## create chromosomes list
       query="SELECT DISTINCT chr FROM knowngene;"
       self.cur.execute(query)
       for e in self.cur:
           self.chrList.append(e[0])
           self.exons[e[0]]=[]
           self.cds[e[0]]=[]
           self.utr[e[0]]=[]
           self.genes[e[0]]=[]
           self.neargenes[e[0]]=[]
       self.plog( "Detected Chromosomes: " + str(len(self.chrList)))

       ## create cache
       self.cur.execute("SELECT id, genelength, exons FROM genelength;")
       for e in self.cur:
           self.geneLength[e[0]]=int(e[1])
           self.geneExons[e[0]]=int(e[2])
       self.plog( "Gene Counted: " + str(len(self.geneExons)))

       ## create cache for exons, gene, UTR, ...
       tempgene=dict()
       self.cur.execute("""SELECT id, chr, chr_start, chr_stop, type FROM knowngene ORDER BY chr, chr_start, chr_stop;""")
       for e in self.cur:
           type = e[4]
           if type == "GENE":
               self.genes[e[1]].append( [ int(e[2]), int(e[3]), e[0] ] )
           if type == "UTR":
               self.utr[e[1]].append( [ int(e[2]), int(e[3]), e[0], '=', '=' ] )
           if type == "CDS":
               self.cds[e[1]].append( [ int(e[2]), int(e[3]), e[0], '=', '=' ] )
           if type == "EXON":
               if e[0] not in tempgene:
                   tempgene[e[0]]=1
               self.exons[e[1]].append( [ int(e[2]), int(e[3]), e[0], tempgene[e[0]], self.geneExons[e[0]] ] )
               tempgene[e[0]]+=1
       for ch in self.chrList:
           self.plog( "Exons in chromosome " + ch + ": " + str(len(self.exons[ch])))
           self.plog( "Genes in chromosome " + ch + ": " + str(len(self.genes[ch])))


   def createCacheFromFile( self, annotations ):
       """Creates the cache needed to get fast query
          Each row must be:
          gene|chr|sense|start|stop|[UTR,CDS,EXON,GENE]
          TXT file was created from database:
          sqlite3 hg19db.sqlite "select distinct annotations.annot, chr,
             strand, chr_start, chr_stop,type from knowngene, annotations
             where annotations.id=knowngene.id order by chr, chr_start,
             chr_stop;"  >  tabella_annotazioni_hg19.txt"""

       for line in open(annotations, "r"):
           e=line.strip().split("|")

           ## create chromosomes list
           if e[1] not in self.chrList:
               self.chrList.append(e[1])
               self.exons[e[1]]=[]
               self.cds[e[1]]=[]
               self.utr[e[1]]=[]
               self.genes[e[1]]=[]

           ## create cache
           if e[5] == "GENE":
               lun= int(e[4]) - int(e[3])
               if e[0] not in self.geneLength or self.geneLength[e[0]]<lun:
                   self.geneLength[e[0]]=int(lun)

           ## create cache for exons, gene, UTR, ...
           tempgene=dict()
           type = e[5]
           if type == "GENE":
               self.genes[e[1]].append( [ int(e[3]), int(e[4]), e[0] ] )
           if type == "UTR":
               self.utr[e[1]].append( [ int(e[3]), int(e[4]), e[0], '=', '=' ] )
           if type == "CDS":
               self.cds[e[1]].append( [ int(e[3]), int(e[4]), e[0], '=', '=' ] )
           if type == "EXON":
               if e[0] not in tempgene:
                   tempgene[e[0]]=1
               self.exons[e[1]].append( [ int(e[3]), int(e[4]), e[0], tempgene[e[0]], "." ] )
               tempgene[e[0]]+=1

       for chr in self.chrList:
           self.cds[chr].sort()
           self.utr[chr].sort()
           self.exons[chr].sort()
           self.genes[chr].sort()


   def createNearGeneCache(self, threshold):
       """Creates the cache needed to get fast query on nearGene coords"""
       self.cur.execute("""SELECT id, chr, chr_start, chr_stop, type FROM knowngene WHERE type='GENE' ORDER BY chr, chr_start, chr_stop;""")
       for e in self.cur:
           ngstart=int(e[2])-threshold
           ngstop =int(e[3])+threshold
           if ngstart<0:
               ngstart=0
           self.neargenes[e[1]].append( [ ngstart, ngstop, e[0] ] )
       for ch in self.chrList:
           self.plog( "Near genes reads in chromosome " + ch + ": " + str(len(self.neargenes[ch])))


   def getDICT(self, whats):
       """whats values are: CDS, UTR, EXON, GENE"""
       if whats=="CDS":
           return self.cds
       if whats=="UTR":
           return self.utr
       if whats=="EXON":
           return self.exons
       if whats=="GENE":
           return self.genes


   def getChrList(self):
       return self.chrList


   def getAnnotationsCached(self, knownGeneID):
       """Returns the gene symbol paired to the given knownGeneID."""
       if knownGeneID in self.splices:
           return self.splices[ knownGeneID ]
       return [ knownGeneID ]


   def getSpliceGeneCached(self, knownGeneID):
       """Returns the gene symbol paired to the given knownGeneID."""
       if knownGeneID in self.splices:
           return ",".join(self.splices[ knownGeneID ])
       return knownGeneID


   def getUTRGeneCached(self, chromosome, sstart, sstop):
       """Returns the list of UTRs containing the given range
          (the range has to be fully covered)."""
       if chromosome in self.utr:
           return binaryFeatSearch( self.utr[chromosome], sstart, sstop, True )
       else:
           return []


   def getCDSGeneCached(self, chromosome, sstart, sstop):
       """Returns the list of CDSs containing the given range
          (the range has to be fully covered)."""
       if chromosome in self.cds:
           return binaryFeatSearch( self.cds[chromosome], sstart, sstop, True )
       else:
           return []


   def getExonGeneCached(self, chromosome, sstart, sstop):
       """Returns the list of exons containing the given range
          (the range has to be fully covered)."""
       if chromosome in self.exons:
           return binaryFeatSearch( self.exons[chromosome], sstart, sstop, True )
       else:
           return []


   def getGeneGeneCached(self, chromosome, sstart, sstop):
       """Returns the list of gene containing the given range
          (the range may be partially covered)."""
       if chromosome in self.genes:
           return binaryFeatSearch( self.genes[chromosome], sstart, sstop, False )
       else:
           return []


   def getNearGeneGeneCached(self, chromosome, sstart, sstop):
       """Returns the list of neargene regions containing the given range
          (the range has to be fully included)."""
       if chromosome in self.neargenes:
           return binaryFeatSearch( self.neargenes[chromosome], sstart, sstop, True )
       else:
           return []


   def getGeneLengthCached(self, gene):
       """Returns the length of the given gene."""
       if gene in geneLength:
           return geneLength[gene]
       return -1


   def geneExonInGeneCached(self, gene):
       """Returns the number of exons (counted by database) of the given gene."""
       if gene in geneExons:
           return geneExons[gene]
       return -1

   ### QUERY DATABASE ###
   ### QUERY DATABASE ###
   ### QUERY DATABASE ###

   def getSpliceGene(self, knownGeneID):
       """Selects gene symbols using the given knownGeneID"""
       geneSymbols=[]
       query="SELECT annot FROM annotations WHERE knowngeneid=?"
       self.cur.execute(query, (knownGeneID,))
       for e in self.cur:
           geneSymbols.append(e[0])
       if len(geneSymbols)==0:
           geneSymbols.append(knownGeneID)
       return ",".join(geneSymbols)


   def getFeaturesGene(self, chromosome, sstart, sstop, type):
       """Selects a type of feature mapping in given position"""
       query="SELECT DISTINCT id FROM knowngene WHERE type=? AND chr=? AND chr_start<=? AND chr_stop>=?;"
       self.cur.execute(query, (type, chromosome, sstart, sstop))
       res=[]
       for e in self.cur:
           res.append(e[0])
       return res

   def getExonGene(self, chromosome, sstart, sstop):
       """Selects the list of exons containing the given range."""
       return self.getFeaturesGene("EXON", chromosome, sstart, sstop)

   def getGeneGene(self, chromosome, sstart, sstop):
       """Selects the list of genes containing the given range."""
       return self.getFeaturesGene("GENE", chromosome, sstart, sstop)

   def getUTRGene(self, chromosome, sstart, sstop):
       """Selects the list of UTR containing the given range."""
       return self.getFeaturesGene("UTR", chromosome, sstart, sstop)

   def getCDSGene(self, chromosome, sstart, sstop):
       """Selects the list of CDS containing the given range."""
       return self.getFeaturesGene("CDS", chromosome, sstart, sstop)


### CLASS COVERAGE ###
### CLASS COVERAGE ###
### CLASS COVERAGE ###

class genomeCoverage:
   oldmerge=dict()
   dest=dict()
   bases=0
   sequences=0

   def __init__(self):
       self.dest=dict()
       oldmerge=dict()
       bases=0


   def getDict(self):
       return self.dest


   def getChromosomes(self):
       chrom=[]
       for ch in self.dest:
           chrom.append( ch )
       chrom.sort()
       return chrom


   def getBases(self):
       return self.bases


   def getSequences(self):
       return self.sequences


   def mergeList ( self, chromosome ):
       new=[]
       for e in self.dest[chromosome]:
           if len(new)==0 or e[0] > new[-1][1] + 1 :
               new.append( e )
           else:
               new[-1] = [ new[-1][0], max( new[-1][1], e[1] ) ]
       self.dest[ chromosome ]=new
       self.oldmerge[ chromosome ]=len(self.dest[ chromosome ])


   def getList( self, chromosome ):
       if chromosome in self.dest:
           if len(self.dest[ chromosome ])!=self.oldmerge[ chromosome ]:
               self.mergeList( chromosome )
           return self.dest[ chromosome ]
       return []


   def addu ( self, element ):
       """ Adds the element to the object. Element is:
           chromosome|start|stop"""
       dd=[]
       el = [ int(element[1]), int(element[2]) ]
       self.bases+=(el[1]-el[0]+1)
       self.sequences+=1

       if element[0] not in self.dest:
           self.dest[ element[0] ]=[ el ]
           self.oldmerge[ element[0] ]=0
           return

       bisect.insort( self.dest[ element[0] ], el )
       if len(self.dest[ element[0] ]) - self.oldmerge[ element[0] ] > 10000:
           self.mergeList( element[0] )


   def addl ( self, elementlist ):
       """ Adds a list of elements to the object. Each element is:
           chromosome|start|stop"""
       for e in elementlist:
           self.addu ( e )


   def getCoverageLength(self):
       ll=0
       for chrom in self.getChromosomes():
           for e in self.getList( chrom ):
               ll+=( 1 + e[1] - e[0] )
       return ll

